/*
 * Copyright 2009-2010 Belmont Software Services
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.belmont.backup;

import java.io.*;
import java.util.*;
import org.json.*;
import java.sql.*;

public class RestoreClient implements IBackupConstants, IEventProducer {
    static final boolean DEBUG = true;
    FileStorage storage;
    File configDir;
    Producer producer;
    DatabaseManifest dbm;
    FileDatabase db;

    public void setDatabase(FileDatabase db) {
	this.db = db;
	this.dbm = new DatabaseManifest(db);
    }

    public void addObserver(IEventObserver observer, int start, int end) {
	producer.addObserver(observer, start, end);
    }
    
    public void removeObserver(IEventObserver observer) {
	producer.removeObserver(observer);
    }

    public RestoreClient(File configDir) {
	this.configDir = configDir;
	if (! configDir.isDirectory()) {
	    configDir.mkdirs();
	}

	this.storage = FileStorage.getStorage(new File(configDir, STORE_DIR).getAbsolutePath());
	producer = Producer.getInstance();
    }

    public boolean restore(String path, File restoreDir, IBackupServer server, String session) throws IOException, InterruptedException, JSONException, SQLException {
	return restore(path, null, restoreDir, server, session);
    }


    int totalRequests;
    int requestsProcessed;

    public boolean restore(String backupname,
			   String path,
			   File restoreDir,
			   IBackupServer server,
			   String session) throws IOException, InterruptedException, JSONException, SQLException {
	return restore(backupname, path, restoreDir, server, session, null);
    }

    public synchronized boolean restore(String backupname,
			   String path,
			   File restoreDir,
			   IBackupServer server,
			   String session,
			   Hashtable filter) throws IOException, InterruptedException, JSONException, SQLException {
	Utils.clearCancel(path);
	if (db == null) {
	    db = new FileDatabase();
	    db.init(new File(configDir, CONFIG_DB_NAME).getAbsolutePath());
	    dbm = new DatabaseManifest(db);
	}

	Status status = Status.getInstance();
	status.reset();
	String abspath = restoreDir.getAbsolutePath();
	Utils.log(LOG_INFO, "Starting restore of "+path+" to "+abspath);

	if (backupname == null) {
	    ResultSet rs = db.getBackupRecords();

	    try  {
		if (rs.next()) {
		    backupname = rs.getString(2);
		} else {
		    Utils.log(LOG_ERROR, "No backup record found");
		    return false;
		}
	    } finally {
		rs.close();
	    }
	}
	Utils.checkCancel();
	producer.sendNotify(this, NOTIFY_RESTORE_START, abspath);
	status.setCurrentTask("Restoring "+abspath);
	//storage.deleteFiles();
	if (! restoreDir.isDirectory()) {
	    restoreDir.mkdirs();
	}


	//producer.sendNotify(this, NOTIFY_SCANNING_START, abspath);
	status.setCurrentOperation("Scanning");
	String restoreid = "restore."+System.currentTimeMillis();

	// create an index of the restore directory (shouldn't we be
	// creating an index of the directory that was backed up since
	// it may have already the files?).
	int total = dbm.createIndex(restoreid, "RestoreClient", null, restoreDir, storage);

	//producer.sendNotify(this, NOTIFY_SCANNING_STOP, abspath);

	ResultSet rs = db.getBackupChecksums(backupname);
	JSONArray files;

	requestsProcessed = 0;
	totalRequests = 0;

	long totalSize = 0;
	total = 0;

	try {
	    while (rs.next()) {
		String ck = rs.getString(1);

		if (filter != null && filter.get(ck) == null) {
		    continue;
		}
		total++;
		totalSize += rs.getLong(2);
	    }
	} finally {
	    rs.close();
	}

	rs = db.getBackupChecksums(backupname);

	//producer.sendNotify(this, NOTIFY_TRANSFER_SIZE, totalSize);
	//producer.sendNotify(this, NOTIFY_FILE_COUNT_CHANGED, new Integer(total));
	status.setAmountToTransfer(totalSize);
	status.setFilesToTransfer(total);
	Utils.log(LOG_INFO,"Restore size "+totalSize+" file count "+total);

	try {
	    while ((files = makeFileRequests(restoreDir, backupname, restoreid, rs, 100, filter)) != null) {
		Utils.checkCancel();
		int processed = processFileRequests(server, session, files);
		if (processed > 0) {
		    copyFiles(backupname, restoreDir, files, filter);
		}
		requestsProcessed += processed;
		totalRequests += processed;
		//producer.sendNotify(this, NOTIFY_COMPLETED_FILE_COUNT_CHANGED, new Integer(totalRequests));
		status.setFilesTransferred(totalRequests);
	    }
	} finally {
	    rs.close();
	}

	try {
	    if (requestsProcessed == totalRequests) {
		producer.sendNotify(this, NOTIFY_RESTORE_DONE, path);
		Utils.log(LOG_INFO, "Restore of "+path+" succesful");
		return true;
	    } else {
		Utils.log(LOG_ERROR, "Error restoring: requestsProcessed = "+requestsProcessed+" totalRequests: "+totalRequests);
		producer.sendNotify(this, NOTIFY_RESTORE_FAILED, path);
		return false;
	    }
	} finally {
	    // Delete all but the last restore records
	    int pathid = db.getPath(abspath);
	    if (pathid != -1) {
		rs = db.getBackupsForPath(pathid);

		while (rs.next()) {
		    String name = rs.getString(2);

		    if (name.equals(restoreid) || !name.startsWith("restore.")) {
			continue;
		    }
		    
		    int rid = rs.getInt(1);

		    db.deleteBackupEntries(rid);
		    db.deleteBackup(rid);
		}
	    }
	}
    }

    /**
     * Iterate over backupentries using rs and create a json array of
     * filerequests for all checksums that are not already present in the
     * restore dir.
     */
    JSONArray makeFileRequests(File restoreDir, String backupName, String restoreid, ResultSet rs, int amount, Hashtable filter) throws IOException, SQLException {
	JSONArray req = new JSONArray();
	int i = 0;
	int rid = db.getBackupId(restoreid);
	Status status = Status.getInstance();

	while (rs.next()) {
	    String checksum = rs.getString(1);
	    long size = rs.getLong(2);

	    if (filter != null && filter.get(checksum) == null) {
		continue;
	    }
	    if (! db.queryChecksum(rid, checksum)) {
		req.put(new FileRequest(checksum, checksum, size, 0));
		if (++i > amount) {
		    return req;
		}
	    } else {
		verifyAllPaths(restoreDir, backupName, checksum);
		totalRequests++;
		requestsProcessed++;
		//producer.sendNotify(this, NOTIFY_COMPLETED_FILE_COUNT_CHANGED, new Integer(totalRequests));
		//producer.sendNotify(this, NOTIFY_TRANSFER_PROGRESS, new Long(size));
		status.setFilesTransferred(totalRequests);
		status.setAmountTransferred(size);
	    }
	}

	return (i == 0) ? null : req;
    }

    /**
     * Verify that all the paths for this checksum exist. If not then
     * copy the ones that exist to the others.
     */
    void verifyAllPaths(File restoreDir, String backupid, String checksum) throws IOException, SQLException {
	ResultSet rs = db.getChecksumPaths(backupid, checksum);
	Vector<String> exists = new Vector<String>();
	Vector<String> failed = null;

	while (rs.next()) {
	    String path = rs.getString(1);
	    File f = new File(restoreDir, path);

	    if (f.exists()) {
		exists.addElement(path);
	    } else {
		if (failed == null) {
		    failed = new Vector<String>();
		}
		failed.addElement(path);
	    }
	}

	if (failed != null) {
	    if (exists.size() == 0) {
		Utils.log(LOG_ERROR, "None of the checksum paths exist for "+checksum);
		return;
	    }
	    File rp = new File(restoreDir, exists.elementAt(0));
	    int l = failed.size();
	    String abspath = rp.getAbsolutePath();
	    for (int i=0; i < l; i++) {
		File np = new File(restoreDir, failed.elementAt(i));
		Utils.log(LOG_INFO, "Copying local file "+abspath+" to "+np.getAbsolutePath());
		try {
		    copyFile(rp, np);
		} catch (IOException ex) {
		    Utils.log(LOG_ERROR, "Copy "+abspath+" "+np.getAbsolutePath(), ex);
		    producer.sendNotify(this, NOTIFY_FILE_ERROR, np.getAbsolutePath());
		}
	    }
	}
    }

    void copyFile(File a, File b) throws IOException {
	FileInputStream astream = new FileInputStream(a);
	byte buf[] = BufferPool.getInstance().get(1024);
	int l;

	try {
	    File p = new File(b.getParent());
	    p.mkdirs();
	    FileOutputStream bstream = new FileOutputStream(b);
	    try {
		while ((l = astream.read(buf)) != -1) {
		    bstream.write(buf, 0, l);
		}
	    } finally {
		BufferPool.getInstance().put(buf);
		bstream.close();
	    }
	} finally {
	    astream.close();
	}
    }

    /**
     * For each checksum in files we use the database to query all of the restore locations
     * and copy the file to each restore location.
     *
     */
    void copyFiles(String backupname, File restoreDir, JSONArray files, Hashtable filter) throws IOException, SQLException, InterruptedException {
	int l = files.length();
	for (int i=0; i < l; i++) {
	    FileRequest r = (FileRequest)files.opt(i);
	    String ck = r.getDigest();
	    ResultSet rs = db.getChecksumPaths(backupname, ck);

	    try {
		Object fpath = (filter == null) ? null : filter.get(ck);
		while (rs.next()) {
		    String path = rs.getString(1);
		    boolean docopy = (fpath == null) || (path.equals(fpath));
		    if (!docopy && fpath != null) {
			if (fpath instanceof Vector) {
			    Vector fv = (Vector)fpath;
			    docopy = fv.contains(path);
			}
		    }

		    Utils.checkCancel();

		    if (docopy) {
			copyFile(restoreDir, ck, path);
		    }
		}
	    } finally {
		SFile sf = storage.get(ck);
		storage.delete(sf);
		rs.close();
	    }
	}
    }

    void copyFile(File restoreDir, String checksum, String path) throws IOException {
	//producer.sendNotify(this, NOTIFY_START_FILE, path);
	Status.getInstance().setCurrentFile(path);
	Status.getInstance().setCurrentOperation("Copying");
	Utils.log(LOG_INFO,"Copy file "+checksum+" to "+restoreDir.toString()+"/"+path);
	SFile sf = storage.get(checksum);
	if (sf == null) {
	    Utils.log(LOG_ERROR,"Error getting storage file for "+path);
	    producer.sendNotify(this, NOTIFY_FILE_ERROR, path);
	    return;
	}

	InputStream in = sf.getInputStream();
	if (in == null) {
	    Utils.log(LOG_ERROR, "Error getting input stream for "+path);
	    producer.sendNotify(this, NOTIFY_FILE_ERROR, path);
	    return;
	}

	byte buf[] = BufferPool.getInstance().get(1024);
	int len;

	File rf = new File(restoreDir, path);

	new File(rf.getParent()).mkdirs();

	try {
	    OutputStream out = new FileOutputStream(rf);

	    try {
		while ((len = in.read(buf)) != -1) {
		    out.write(buf, 0, len);
		}
	    } finally {
		out.close();
	    }
	} finally {
	    BufferPool.getInstance().put(buf);
	    in.close();
	}
	Utils.log(LOG_INFO, "File copied "+checksum);
	//producer.sendNotify(this, NOTIFY_FILE_DONE, path);
    }

    int processFileRequests(IBackupServer server, String session, JSONArray requests) throws InterruptedException {
	int requestsSent = 0;
	int l = requests.length();

	Status status = Status.getInstance();
	status.setCurrentOperation("Downloading");
	for (int i=0; i < l; i++) {
	    JSONObject request = (JSONObject)requests.opt(i);
	    String name = FileRequest.getName(request);
	    String digest = FileRequest.getDigest(request);
	    long offset = FileRequest.getOffset(request);
	    long size = FileRequest.getLength(request);

	    Utils.checkCancel();
	    //producer.sendNotify(this, NOTIFY_START_FILE, name);
	    status.setCurrentFile(name);
	    try {
		SFile sf = storage.getFileCreate(digest);

		if (sf == null) {
		    Utils.log(LOG_ERROR,"Error finding file for "+digest);
		    producer.sendNotify(this, NOTIFY_FILE_ERROR, name);
		    continue;
		}

		OutputStream out = null;

		try {
		    out = sf.getOutputStream(offset);
		} catch (IOException ex) {
		    // try again requesting the whole file
		    if (offset > 0) {
			Utils.log(LOG_INFO,"Retrying getOutputStream for "+digest+" offset "+offset);
			offset = 0;
			out = sf.getOutputStream(offset);
		    }
		}

		if (out == null) {
		    Utils.log(LOG_ERROR,"Error getting output stream for "+digest);
		    producer.sendNotify(this, NOTIFY_FILE_ERROR, name);
		} else {
		    try {
			server.getFile(session, digest, name, size, offset, out);
			requestsSent++;
			//producer.sendNotify(this, NOTIFY_FILE_DONE, name);
		    } finally {
			out.close();
		    }
		    sf.commit();
		}
	    } catch (IOException ex) {
		producer.sendNotify(this, NOTIFY_FILE_ERROR, name);
		Utils.log(LOG_ERROR, "File exception for "+name, ex);
	    }
	}

	return requestsSent;
    }

    public static void main(String args[]) throws Exception {
	RestoreClient client = new RestoreClient(new File(args[2]));
	FileBackupServer server = new FileBackupServer(new File(args[3]));
	String session = server.connect(BACKUP_PROTOCOL_VERSION, "test.backupclient", "localhost",
					System.getProperty("os.name"),
					System.getProperty("os.version"),
					"BackupClientUser");
	if (client.restore(null, args[0], new File(args[1]), server, session)) {
	    System.out.println("restore done");
	} else {
	    System.out.println("restore failed");
	}
	server.disconnect(session);
	System.exit(0);
    }
}